… pool (meta-pytorch#3533)

Summary:

A few changes in the diff:
1. Support to proportionally shard the tensor pool based on memory capacity per rank. 
2. Using block_bucketize_sparse_features_inference to return bucket_mapping that can be used during request batching in inference w/ custom sigrid predictor engine
3. Wrapping some of the operations with fx wrappers to make it compatible with model split boundaries for DLRM serving where embeddings are sharded and split onto different pytorch modules
4. Exposing set_device() api to some of the modules if we want to place some shards to cpu while others to cuda.
5. Move _get_unbucketize_tensor_via_length_alignment to common util files.

Differential Revision: D79603009

… pool (meta-pytorch#3533)

Summary:

A few changes in the diff:
1. Support to proportionally shard the tensor pool based on memory capacity per rank. 
2. Using block_bucketize_sparse_features_inference to return bucket_mapping that can be used during request batching in inference w/ custom sigrid predictor engine
3. Wrapping some of the operations with fx wrappers to make it compatible with model split boundaries for DLRM serving where embeddings are sharded and split onto different pytorch modules
4. Exposing set_device() api to some of the modules if we want to place some shards to cpu while others to cuda.
5. Move _get_unbucketize_tensor_via_length_alignment to common util files.

As part of this diff, also had to update some of the test cases. Mainly because updating the forward path a bit leads to reorganization of return values of remote module in the test cases leading to reorganization of batchinfo for each of those output. 

Baseline test
Full Output: https://www.internalfb.com/intern/everpaste/?handle=GMpN0B-VwcKcvAgDAB_hIkOvKrJNbsIXAAAB&phabricator_paste_number=2035160680
Remote graph: https://www.internalfb.com/intern/everpaste/?color=0&handle=GDe4BCBV-XwTNLoEAOTnVbCneb1jbr0LAAAz
Output order: (_item_embedding_index_values_tensor_pool__local_shard_pools_0, _item_embedding_index_values_tensor_pool__local_shard_pools_1, getitem_6, getitem_10, getitem_9)

After changes
Full Output: https://www.internalfb.com/intern/everpaste/?handle=GIb_pSL6TjHawBgEALDgarUck4YhbsIXAAAB&phabricator_paste_number=2035191658 
Remote graph: https://www.internalfb.com/intern/everpaste/?color=0&handle=GFmFkB9waX3elzMGAJr9zTcLZiIDbr0LAAAz
Output Order: getitem_6, _item_embedding_index_values_tensor_pool__local_shard_pools_0, _item_embedding_index_values_tensor_pool__local_shard_pools_1, getitem_10, getitem_9

getitem_6 is shifted first after changes.

Differential Revision: D79603009